package asm;

import org.objectweb.asm.*;
import org.objectweb.asm.commons.AdviceAdapter;
import org.objectweb.asm.commons.Method;
import org.objectweb.asm.tree.AnnotationNode;

import java.io.FileInputStream;
import java.io.FileOutputStream;

public class ASMLoadDemo {

    public static void main(String[] args) throws Exception {
        byte[] bytes = insertCodeByLoadClass("D:\\20210426\\code\\JavaProjec\\JMMDemo\\out\\production\\JMMDemo\\asm\\InjectTest.class");

        //测试
        ASMGenDemo.testClass("asm.InjectTest",bytes);
    }

    /**
     * 加载Class插入代码
     *
     * @param path
     * @return
     */
    public static byte[] insertCodeByLoadClass(String path) {
        try {
            //1.获取Class文件
            FileInputStream fis = new FileInputStream(path);
            //2.读取Class文件
            ClassReader classReader = new ClassReader(fis);
            //3.准备
            ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_FRAMES);
            //4.重点！！ 开始插桩 对代码进行扫描，进行代码插入。
            classReader.accept(new MyClassVisitor(Opcodes.ASM9, classWriter), ClassReader.SKIP_DEBUG);
            //5.插桩结束 输出字节码
            byte[] bytes = classWriter.toByteArray();
            //6.保存文件
            FileOutputStream fos = new FileOutputStream("src/asm/InjectTest.class");
            fos.write(bytes);
            fos.close();
            fis.close();
            return bytes;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 用来观察（扫描）类信息，包括变量、方法
     */
    static class MyClassVisitor extends ClassVisitor {

        public MyClassVisitor(int api) {
            super(api);
        }

        public MyClassVisitor(int api, ClassVisitor classVisitor) {
            super(api, classVisitor);
        }

        /**
         * 当有一个方法，就执行这个回调一次，类中有多个方法，这里就会执行多次
         */
        @Override
        public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
            System.out.println("读到Class的一个方法visitMethod: " + name);
//            if (!name.equals("test")) {
//                return super.visitMethod(access, name, descriptor, signature, exceptions);
//            }
            MethodVisitor methodVisitor = super.visitMethod(access, name, descriptor, signature, exceptions);
            return new MyMethodVisitor(api, methodVisitor, access, name, descriptor);
        }


        /**
         * 当有一个变量，就执行这个回调一次，类中有多个变量，这里就会执行多次
         */
        @Override
        public FieldVisitor visitField(int access, String name, String descriptor, String signature, Object value) {
            System.out.println("读到Class的一个变量visitField: " + name);
//            if (name.equals("a")) {
//                //删除a变量
//                return null;
//            }
//            if (name.equals("b")) {
//                //修改b变量名词为d，值为10
//                return super.visitField(access, name, descriptor, signature, 10);
//            }
            FieldVisitor fieldVisitor = super.visitField(access, name, descriptor, signature, value);
            return new MyFieldVisitor(api, fieldVisitor);
        }

        /**
         * @param descriptor
         * @param visible
         * @return
         */
        @Override
        public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
            System.out.println("读到Class的一个注解visitAnnotation: " + descriptor);
            AnnotationVisitor annotationVisitor = super.visitAnnotation(descriptor, visible);

            MyAnnotationVisitor myAnnotationVisitor = new MyAnnotationVisitor(api, descriptor);
            myAnnotationVisitor.accept(annotationVisitor);
            return myAnnotationVisitor;
        }

        @Override
        public void visitEnd() {
            super.visitEnd();
            //添加一个变量
            FieldVisitor fv = cv.visitField(Opcodes.ACC_PRIVATE, "e", "I", null, null);
            if (fv != null) {
                fv.visitEnd();
            }
        }
    }

    /**
     * 用来观察（扫描）方法里的代码
     */
    static class MyMethodVisitor extends AdviceAdapter {

        protected MyMethodVisitor(int api, MethodVisitor methodVisitor, int access, String name, String descriptor) {
            super(api, methodVisitor, access, name, descriptor);
        }

        int startTimeIndex;//startTime在局部变量表的位置
        int endTimeIndex;//endTime在局部变量表的位置
        boolean needInsertCode = false;//标记是否需要插入代码

        /**
         * 方法进入的时候调用
         */
        @Override
        protected void onMethodEnter() {
            super.onMethodEnter();
            if (!needInsertCode) {
                return;
            }

            //调用System.currentTimeMillis方法，调用完后将返回值推到栈顶
            invokeStatic(Type.getType("Ljava/lang/System;"), new Method("currentTimeMillis", "()J"));
            //获取当前局部变量表可插入的位置，不能乱传，不然会影响下一个值的索引获取，这里你传DOUBLE_TYPE或者LONG_TYPE,都是一样的
            startTimeIndex = newLocal(Type.LONG_TYPE);
            //把栈顶的值（也就是System.currentTimeMillis的返回值）保存到局部变量表startTimeIndex位置
            storeLocal(startTimeIndex);
            System.out.println("onMethodEnter getName: " + getName() + " startIndex " + startTimeIndex);
        }

        @Override
        public void visitLocalVariable(String name, String descriptor, String signature, Label start, Label end, int index) {
            System.out.println("visitLocalVariable>>>>>>>>>>>>>>>>>>>>>>");
            super.visitLocalVariable(name, descriptor, signature, start, end, index);
        }

        @Override
        public void visitFrame(int type, int numLocal, Object[] local, int numStack, Object[] stack) {
            System.out.println("visitFrame>>>>>>>>>>>>>>>>>>>>>>");

            super.visitFrame(type, numLocal, local, numStack, stack);
        }

        /**
         * 当方法退出的时候调用
         */
        @Override
        protected void onMethodExit(int opcode) {
            super.onMethodExit(opcode);
            if (!needInsertCode) {
                return;
            }
            //调用System.currentTimeMillis方法，调用完后将返回值推到栈顶
            invokeStatic(Type.getType("Ljava/lang/System;"), new Method("currentTimeMillis", "()J"));
            //获取当前局部变量表可插入的位置
            endTimeIndex = newLocal(Type.LONG_TYPE);
            //把栈顶的值（也就是System.currentTimeMillis的返回值）保存到局部变量表endTimeIndex位置
            storeLocal(endTimeIndex);
            //调用System.out，调用完后栈顶会有一个
            getStatic(Type.getType("Ljava/lang/System;"), "out", Type.getType("Ljava/io/PrintStream;"));
            //new 一个 StringBuilder();并将其引用值压入栈顶
            newInstance(Type.getType("Ljava/lang/StringBuilder;"));
            //复制栈顶数值并将复制值压入栈顶（对象StringBuilder的引用）
            dup();
            //调用StringBuilder的构造方法
            invokeConstructor(Type.getType("Ljava/lang/StringBuilder;"), new Method("<init>", "()V"));
            //将"execute time=" 从常量池中推送至栈顶
            visitLdcInsn("execute time=");
            //调用StringBuilder的append方法，将栈顶的值添加进StringBuilder
            invokeVirtual(Type.getType("Ljava/lang/StringBuilder;"), new Method("append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;"));
            //把局部变量表中的endTimeIndex位置元素，压入栈顶
            loadLocal(endTimeIndex);
            //把局部变量表中的startTimeIndex位置元素，压入栈顶
            loadLocal(startTimeIndex);
            //将栈顶两元素进行相减，把结果再压入栈顶
            math(SUB, Type.LONG_TYPE);
            //调用append方法，将栈顶值添加
            invokeVirtual(Type.getType("Ljava/lang/StringBuilder;"), new Method("append", "(J)Ljava/lang/StringBuilder;"));
            //将" ms" 从常量池中推送至栈顶
            visitLdcInsn(" ms");
            ////调用StringBuilder的append方法，将栈顶的值添加进StringBuilder
            invokeVirtual(Type.getType("Ljava/lang/StringBuilder;"), new Method("append", "(Ljava/lang/String;)Ljava/lang/StringBuilder;"));
            ////调用StringBuilder的toString方法,调用完后，将结果推到栈顶
            invokeVirtual(Type.getType("Ljava/lang/StringBuilder;"), new Method("toString", "()Ljava/lang/String;"));
            //调用println，打印栈顶的值
            invokeVirtual(Type.getType("Ljava/io/PrintStream;"), new Method("println", "(Ljava/lang/String;)V"));
        }

        @Override
        public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
            if ("Lasm/ASMTarget;".equals(descriptor)) {
                needInsertCode = true;
            }
            return super.visitAnnotation(descriptor, visible);
        }
    }

    static class MyFieldVisitor extends FieldVisitor {

        protected MyFieldVisitor(int api) {
            super(api);
        }

        protected MyFieldVisitor(int api, FieldVisitor fieldVisitor) {
            super(api, fieldVisitor);
        }

        @Override
        public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
            System.out.println("MyFieldVisitor 读到一个变量上的注解 descriptor：" + descriptor + " visible: " + visible);
            return super.visitAnnotation(descriptor, visible);
        }

        @Override
        public AnnotationVisitor visitTypeAnnotation(int typeRef, TypePath typePath, String descriptor, boolean visible) {
            System.out.println("MyFieldVisitor 读到一个变量上的注解类型 typeRef： " + typeRef + " descriptor: " + descriptor);
            return super.visitTypeAnnotation(typeRef, typePath, descriptor, visible);
        }

        @Override
        public void visitAttribute(Attribute attribute) {
            System.out.println("MyFieldVisitor 读到一个变量上的属性 attribute：" + attribute.type);
            super.visitAttribute(attribute);
        }

        @Override
        public void visitEnd() {
            super.visitEnd();
        }
    }

    static class MyAnnotationVisitor extends AnnotationNode {


        public MyAnnotationVisitor(String descriptor) {
            super(descriptor);
        }

        public MyAnnotationVisitor(int api, String descriptor) {
            super(api, descriptor);
        }

        @Override
        public AnnotationVisitor visitAnnotation(String name, String descriptor) {
            System.out.println("MyAnnotationVisitor visitAnnotation name：" + name + " descriptor: " + descriptor);
            return super.visitAnnotation(name, descriptor);
        }

        @Override
        public AnnotationVisitor visitArray(String name) {
            System.out.println("MyAnnotationVisitor visitArray ");

            return super.visitArray(name);
        }

        @Override
        public void visitEnum(String name, String descriptor, String value) {
            System.out.println("MyAnnotationVisitor visitEnum ");

            super.visitEnum(name, descriptor, value);
        }


        @Override
        public void visitEnd() {
            super.visitEnd();
            System.out.println("MyAnnotationVisitor visitEnd ");

        }
    }

}
